import pandas as pd
import numpy as np
import os
import joblib
from sklearn.model_selection import train_test_split
from sklearn.metrics import accuracy_score, classification_report, roc_auc_score, confusion_matrix, roc_curve
from sklearn.preprocessing import StandardScaler
import sys
import matplotlib.pyplot as plt

# 添加当前目录到系统路径，以便导入model_factory和B1FeatureExtractor
sys.path.append(os.path.dirname(os.path.abspath(__file__)))
from model_factory import ModelFactory
from B1_train import B1FeatureExtractor

def load_all_stock_data(data_dir, max_stocks=None):
    """
    加载所有股票数据（带日期过滤）
    
    Parameters:
    data_dir: str, 数据目录路径
    max_stocks: int, 最大加载股票数量，None表示加载所有
    
    Returns:
    list: 所有股票数据的列表
    """
    stock_files = [f for f in os.listdir(data_dir) if f.endswith('.csv')]
    
    if max_stocks:
        stock_files = stock_files[:max_stocks]
    
    all_data = []
    for file in stock_files:
        file_path = os.path.join(data_dir, file)
        try:
            df = pd.read_csv(file_path)
            # 特征列验证
            required_columns = ['j', 'Vol_x_Price_Change', 'Relative_Volume', 'ATR',
                               'ATR_percentage', 'MA_Trend_Signal', 'Close_Change_Rate', 'Volume_Change_Rate', 'dif']
            missing_cols = [col for col in required_columns if col not in df.columns]
            if missing_cols:
                print(f"警告：文件 {file} 缺失必要列: {missing_cols}")
                continue
            
            df['date'] = pd.to_datetime(df['date'])
            df = df.sort_values('date').reset_index(drop=True)
            
            # 过滤数据，只保留2021年12月1日 - 2024年12月30日的数据
            start_date = pd.to_datetime('2021-12-01')
            end_date = pd.to_datetime('2024-12-30')
            df = df[(df['date'] >= start_date) & (df['date'] <= end_date)].copy()
            
            all_data.append(df)
        except Exception as e:
            print(f"加载文件 {file} 时出错: {e}")
    
    print(f"成功加载 {len(all_data)} 只股票数据")
    return all_data

def load_all_stock_data_no_date_filter(data_dir, max_stocks=None):
    """
    加载所有股票数据（不带日期过滤，用于验证）
    
    Parameters:
    data_dir: str, 数据目录路径
    max_stocks: int, 最大加载股票数量，None表示加载所有
    
    Returns:
    list: 所有股票数据的列表
    """
    stock_files = [f for f in os.listdir(data_dir) if f.endswith('.csv')]
    
    if max_stocks:
        stock_files = stock_files[:max_stocks]
    
    all_data = []
    for file in stock_files:
        file_path = os.path.join(data_dir, file)
        try:
            df = pd.read_csv(file_path)
            # 特征列验证
            required_columns = ['j', 'Vol_x_Price_Change', 'Relative_Volume', 'ATR',
                               'ATR_percentage', 'MA_Trend_Signal', 'Close_Change_Rate', 'Volume_Change_Rate', 'dif']
            missing_cols = [col for col in required_columns if col not in df.columns]
            if missing_cols:
                print(f"警告：文件 {file} 缺失必要列: {missing_cols}")
                continue
            
            # 确保日期列被正确转换为datetime类型
            df['date'] = pd.to_datetime(df['date'])
            df = df.sort_values('date').reset_index(drop=True)
            
            # 不进行日期过滤，保留所有数据用于验证
            all_data.append(df)
        except Exception as e:
            print(f"加载文件 {file} 时出错: {e}")
    
    print(f"成功加载 {len(all_data)} 只股票数据（无日期过滤）")
    return all_data

def load_all_stock_data_2025(data_dir, max_stocks=None):
    """
    加载所有股票数据（专门加载2025年数据用于验证，截止到2025年7月30日）
    
    Parameters:
    data_dir: str, 数据目录路径
    max_stocks: int, 最大加载股票数量，None表示加载所有
    
    Returns:
    list: 所有股票数据的列表
    """
    stock_files = [f for f in os.listdir(data_dir) if f.endswith('.csv')]
    
    if max_stocks:
        stock_files = stock_files[:max_stocks]
    
    all_data = []
    for file in stock_files:
        file_path = os.path.join(data_dir, file)
        try:
            df = pd.read_csv(file_path)
            # 特征列验证
            required_columns = ['j', 'Vol_x_Price_Change', 'Relative_Volume', 'ATR',
                               'ATR_percentage', 'MA_Trend_Signal', 'Close_Change_Rate', 'Volume_Change_Rate', 'dif']
            missing_cols = [col for col in required_columns if col not in df.columns]
            if missing_cols:
                print(f"警告：文件 {file} 缺失必要列: {missing_cols}")
                continue
            
            # 确保日期列被正确转换为datetime类型
            df['date'] = pd.to_datetime(df['date'])
            df = df.sort_values('date').reset_index(drop=True)
            
            # 只保留2025年的数据用于验证（截止到2025年7月30日）
            start_date = pd.to_datetime('2025-01-01')
            end_date = pd.to_datetime('2025-07-30')
            df = df[(df['date'] >= start_date) & (df['date'] <= end_date)].copy()
            
            if not df.empty:
                all_data.append(df)
        except Exception as e:
            print(f"加载文件 {file} 时出错: {e}")
    
    print(f"成功加载 {len(all_data)} 只股票数据（2025年1月1日至7月30日数据）")
    return all_data

def prepare_all_data(all_data, sequence_length=15):
    """
    准备所有股票的数据用于训练
    
    Parameters:
    all_data: list, 所有股票数据
    sequence_length: int, 序列长度
    
    Returns:
    X: 特征矩阵
    y: 标签向量
    feature_extractor: 特征提取器对象（包含scaler）
    """
    feature_extractor = B1FeatureExtractor(sequence_length=sequence_length)
    
    all_X = []
    all_y = []
    
    for i, df in enumerate(all_data):
        try:
            X, y = feature_extractor.prepare_features(df, target_col='B1_tag')
            if len(X) > 0 and len(y) > 0:
                # 只保留J值小于10的样本
                y_binary = (y == 1).astype(int)
                # 特征精度控制
                X = np.array(X).round(4)  # 符合特征工程规范13.1
                all_X.append(X)
                all_y.append(y_binary)
                
            if (i + 1) % 100 == 0:
                print(f"已处理 {i + 1} 只股票的数据")
        except Exception as e:
            print(f"处理第 {i + 1} 只股票时出错: {e}")
    
    if all_X:
        X_combined = np.vstack(all_X)
        y_combined = np.hstack(all_y)
        # 缺失值处理增强
        if np.isnan(X_combined).any():
            print(f"警告：发现特征数据包含NaN值，使用前向填充和后向填充处理")
            X_combined = pd.DataFrame(X_combined).ffill().bfill().fillna(0).values  # 符合数据处理规范8.1
        print(f"合并后的数据形状: X={X_combined.shape}, y={y_combined.shape}")
        return X_combined, y_combined, feature_extractor
    else:
        print("没有足够的数据用于训练")
        return np.array([]), np.array([]), feature_extractor

def plot_roc_curve(y_true, y_pred_proba, filename='roc_curve.png'):
    """
    绘制ROC曲线
    
    Parameters:
    y_true: 真实标签
    y_pred_proba: 预测概率
    filename: 保存的文件名
    """
    fpr, tpr, _ = roc_curve(y_true, y_pred_proba)
    auc_score = roc_auc_score(y_true, y_pred_proba)
    
    plt.figure(figsize=(8, 6))
    plt.plot(fpr, tpr, color='darkorange', lw=2, label=f'ROC curve (AUC = {auc_score:.2f})')
    plt.plot([0, 1], [0, 1], color='navy', lw=2, linestyle='--')
    plt.xlim([0.0, 1.0])
    plt.ylim([0.0, 1.05])
    plt.xlabel('False Positive Rate')
    plt.ylabel('True Positive Rate')
    plt.title('Receiver Operating Characteristic')
    plt.legend(loc="lower right")
    plt.grid(True)
    # 保存ROC曲线图像
    plt.savefig(filename)
    plt.close()

def train_model_with_all_data():
    """
    使用100只股票数据训练模型
    使用2025年股票数据进行验证
    在保留足够准确度情况下，减小模型大小
    可自行调整训练和验证数据数量
    """
    # 数据目录
    data_dir = os.path.join("..", "tdx_data_process")
    
    # 检查数据目录是否存在
    if not os.path.exists(data_dir):
        print(f"数据目录不存在: {data_dir}")
        return
    
    # 加载训练股票数据，限制为前100个股票
    print("开始加载前100个股票数据用于训练...")
    train_data = load_all_stock_data(data_dir, max_stocks=100)
    
    if not train_data:
        print("没有加载到训练数据")
        return
    
    # 准备训练数据
    print("准备训练数据...")
    X_train, y_train, feature_extractor = prepare_all_data(train_data)
    
    if len(X_train) == 0 or len(y_train) == 0:
        print("没有足够的训练数据")
        return
    
    # 显示训练数据分布
    print(f"训练数据分布: B1_tag=1 的样本数: {np.sum(y_train == 1)}, B1_tag!=1 的样本数: {np.sum(y_train != 1)}")
    
    # 加载2025年股票数据用于验证
    print("开始加载2025年股票数据用于验证...")
    all_data = load_all_stock_data_2025(data_dir, max_stocks=3000)
    
    if not all_data:
        print("没有加载到2025年验证数据，尝试使用所有数据进行验证...")
        all_data = load_all_stock_data_no_date_filter(data_dir, max_stocks=500)
    
    if not all_data:
        print("没有加载到验证数据")
        return
    
    # 准备验证数据
    print("准备验证数据...")
    X_test, y_test, _ = prepare_all_data(all_data)
    
    if len(X_test) == 0 or len(y_test) == 0:
        print("没有足够的验证数据")
        return
    
    # 显示验证数据分布
    print(f"验证数据分布: B1_tag=1 的样本数: {np.sum(y_test == 1)}, B1_tag!=1 的样本数: {np.sum(y_test != 1)}")
    
    # 创建并训练多种类型的模型
    model_types = ['rf', 'lr', 'svm','nn']  # 四种模型类型：随机森林、逻辑回归、支持向量机、神经网络
    results = {}
    
    for model_type in model_types:
        print(f"\n{'='*50}")
        print(f"训练 {model_type.upper()} 模型...")
        print(f"{'='*50}")
        
        # 创建模型
        if model_type == 'rf':
            # 随机森林模型
            model = ModelFactory.create_model('rf', n_estimators=100, random_state=42)
        elif model_type == 'lr':
            # 逻辑回归模型
            model = ModelFactory.create_model('lr', random_state=42)
        elif model_type == 'svm':
            # 支持向量机模型
            model = ModelFactory.create_model('svm', random_state=42)
        elif model_type == 'nn':
            # 神经网络模型
            model = ModelFactory.create_model('nn', hidden_layer_sizes=(64, 32), max_iter=500, random_state=42)
        else:
            raise ValueError(f"不支持的模型类型: {model_type}")
        
        # 训练模型
        model.train(X_train, y_train)
        print(f"{model_type.upper()}模型已训练完毕 ")
        
        # 在测试集上进行预测
        y_pred_proba = model.predict(X_test)
        
        # 评估模型
        accuracy = accuracy_score(y_test, (y_pred_proba > 0.5).astype(int))
        try:
            auc_score = roc_auc_score(y_test, y_pred_proba)
        except:
            auc_score = 0.0
        
        print(f"模型评估结果 (使用2025年数据验证):")
        print(f"准确率: {accuracy:.4f}")
        print(f"AUC Score: {auc_score:.4f}")
        
        results[model_type] = {'accuracy': accuracy, 'auc': auc_score}
        
        # 保存模型
        model_path = os.path.join("b1_model", f"b1_{model_type}_all_stocks_model.pkl")
        model.save_model(model_path)
        print(f"模型已保存到: {model_path}")
    
    # 保存scaler
    scaler_path = os.path.join("b1_model", "b1_all_stocks_scaler.pkl")
    joblib.dump(feature_extractor.scaler, scaler_path)
    print(f"scaler已保存到: {scaler_path}")
    
    # 输出模型比较结果
    print(f"\n{'='*50}")
    print("模型比较结果:")
    print(f"{'='*50}")
    for model_type, metrics in results.items():
        print(f"{model_type.upper()}: 准确率={metrics['accuracy']:.4f}, AUC={metrics['auc']:.4f}")

def train_best_models():
    """
    训练四个模型的最佳效果版本，并保存到b1_model_best目录
    根据测试记录：
    - RF和NN使用3000训练集最佳
    - LR使用1000训练集最佳
    - SVM使用100训练集最佳
    """
    # 数据目录
    data_dir = os.path.join("..", "tdx_data_process")
    
    # 检查数据目录是否存在
    if not os.path.exists(data_dir):
        print(f"数据目录不存在: {data_dir}")
        return
    
    # 创建保存最佳模型的目录
    best_model_dir = "b1_model_best"
    os.makedirs(best_model_dir, exist_ok=True)
    
    # 定义各模型的最佳配置
    model_configs = {
        'rf': {'train_stocks': 3000, 'params': {'n_estimators': 100, 'random_state': 42}},
        'lr': {'train_stocks': 1000, 'params': {'random_state': 42}},
        'svm': {'train_stocks': 100, 'params': {'random_state': 42}},
        'nn': {'train_stocks': 3000, 'params': {'hidden_layer_sizes': (64, 32), 'max_iter': 500, 'random_state': 42}}
    }
    
    # 加载2025年股票数据用于验证
    print("开始加载2025年股票数据用于验证...")
    test_data = load_all_stock_data_2025(data_dir, max_stocks=3000)
    
    if not test_data:
        print("没有加载到2025年验证数据，尝试使用所有数据进行验证...")
        test_data = load_all_stock_data_no_date_filter(data_dir, max_stocks=500)
    
    if not test_data:
        print("没有加载到验证数据")
        return
    
    # 准备验证数据
    print("准备验证数据...")
    X_test, y_test, _ = prepare_all_data(test_data)
    
    if len(X_test) == 0 or len(y_test) == 0:
        print("没有足够的验证数据")
        return
    
    # 显示验证数据分布
    print(f"验证数据分布: B1_tag=1 的样本数: {np.sum(y_test == 1)}, B1_tag!=1 的样本数: {np.sum(y_test != 1)}")
    
    results = {}
    scalers = {}  # 保存每个模型对应的scaler
    
    # 逐个训练各模型的最佳版本
    for model_type, config in model_configs.items():
        print(f"\n{'='*50}")
        print(f"训练 {model_type.upper()} 最佳模型 (使用{config['train_stocks']}只股票训练)...")
        print(f"{'='*50}")
        
        # 加载指定数量的训练数据，使用相同的时间范围过滤（2021-12-01 到 2024-12-30）
        print(f"开始加载{config['train_stocks']}只股票数据用于训练...")
        train_data = load_all_stock_data(data_dir, max_stocks=config['train_stocks'])
        
        if not train_data:
            print(f"没有加载到{config['train_stocks']}只训练数据")
            continue
        
        # 准备训练数据
        print("准备训练数据...")
        X_train, y_train, feature_extractor = prepare_all_data(train_data)
        
        if len(X_train) == 0 or len(y_train) == 0:
            print("没有足够的训练数据")
            continue
        
        # 显示训练数据分布
        print(f"训练数据分布: B1_tag=1 的样本数: {np.sum(y_train == 1)}, B1_tag!=1 的样本数: {np.sum(y_train != 1)}")
        
        # 保存该模型对应的scaler
        scalers[model_type] = feature_extractor.scaler
        
        # 创建模型
        model = ModelFactory.create_model(model_type, **config['params'])
        
        # 训练模型
        model.train(X_train, y_train)
        print(f"{model_type.upper()}最佳模型已训练完毕 ")
        
        # 在测试集上进行预测
        y_pred_proba = model.predict(X_test)
        
        # 评估模型
        accuracy = accuracy_score(y_test, (y_pred_proba > 0.5).astype(int))
        try:
            auc_score = roc_auc_score(y_test, y_pred_proba)
        except:
            auc_score = 0.0
        
        print(f"模型评估结果 (使用2025年数据验证):")
        print(f"准确率: {accuracy:.4f}")
        print(f"AUC Score: {auc_score:.4f}")
        
        results[model_type] = {'accuracy': accuracy, 'auc': auc_score}
        
        # 保存模型到b1_model_best目录
        model_path = os.path.join(best_model_dir, f"b1_{model_type}_best_model.pkl")
        model.save_model(model_path)
        print(f"模型已保存到: {model_path}")
        
        # 保存该模型对应的scaler
        scaler_path = os.path.join(best_model_dir, f"b1_{model_type}_best_scaler.pkl")
        joblib.dump(feature_extractor.scaler, scaler_path)
        print(f"scaler已保存到: {scaler_path}")
    
    # 输出模型比较结果
    print(f"\n{'='*50}")
    print("最佳模型比较结果:")
    print(f"{'='*50}")
    for model_type, metrics in results.items():
        print(f"{model_type.upper()}: 准确率={metrics['accuracy']:.4f}, AUC={metrics['auc']:.4f}")

def main():
    """
    主函数
    """
    import argparse
    
    parser = argparse.ArgumentParser(description='训练B1模型')
    parser.add_argument('--mode', type=str, default='standard', 
                        choices=['standard', 'best'],
                        help='训练模式: standard(标准模式) 或 best(最佳模型模式)')
    
    args = parser.parse_args()
    
    if args.mode == 'standard':
        print("使用前100个股票数据训练B1模型，使用全部股票数据进行验证")
        print("=" * 50)
        train_model_with_all_data()
    elif args.mode == 'best':
        print("训练各模型的最佳版本")
        print("=" * 50)
        train_best_models()

if __name__ == "__main__":
    main()